#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
四足机器人轨迹规划程序主模块
简化版本，使用配置文件控制，可选Tk界面
"""

import os
import sys
import time
import threading
import numpy as np
import pybullet as p
import csv
# 导入自定义模块
from config import (
    URDF_FILE, SIMULATION_STEP_TIME, MOVEMENT_SPEED, 
    USE_TK_INTERFACE
)
from four_legged_robot import FourLeggedRobot
from trajectory_planner import TrajectoryPlanner
from trajectory_visualizer import TrajectoryVisualizer

csv_file = open('gait_data.csv', 'w', newline='')
csv_writer = csv.writer(csv_file)
csv_writer.writerow(['timestamp', 'time_phase', 'leg1_status', 'leg2_status', 'leg3_status', 'leg4_status','sum_status'])


# 如果启用Tk界面，导入相关模块
if USE_TK_INTERFACE:
    from tk_interface import TkInterface

# 全局变量
running = True  # 控制仿真运行状态
start_time = time.time()
def simulation_thread(robot, planner, visualizer, ui=None):
    """仿真线程
    
    参数:
        robot: FourLeggedRobot实例
        planner: TrajectoryPlanner实例
        visualizer: TrajectoryVisualizer实例
        ui: TkInterface实例 (可选)
    """
    global running
    
    # 当前角度字典
    current_angles = {}
    for leg_id in robot.leg_chains.keys():
        current_angles[leg_id] = np.zeros(len(robot.leg_chains[leg_id].links))
    
    # 主循环
    while True:
        try:
            # 检查界面是否关闭
            if ui and hasattr(ui, 'root') and not ui.root.winfo_exists():
                break
            
            # 如果使用Tk界面，更新规划器参数
            if ui:
                ui.update_planner_params()
  
            # 如果运行标志为真，执行轨迹规划
            if running:
                # 获取当前时间相位
                current_time_phase = time.time() #速度 更改
                print(f"当前时间相位: {current_time_phase:.2f}")
                current_time_phase = ((current_time_phase) % 5) / 5  #速度 更改
                all_status = []
                # 为每条腿生成轨迹
                for leg_id in robot.leg_chains.keys():
                    # 生成目标位置
                    target_pos, status = planner.generate_trajectory(leg_id, current_time_phase)
                    all_status.append(status)
                    # 求解IK并应用
                    solution, error = robot.solve_ik_and_apply(
                        leg_id, 
                        target_pos, 
                        init_pos=current_angles[leg_id]
                    )
                    
                    # 更新当前角度
                    current_angles[leg_id] = solution
                    
                    # 更新标记位置
                    world_pos = robot.get_world_position(target_pos)
                    visualizer.update_marker_position(leg_id, world_pos)
                print(" ".join(map(str, all_status)))
                csv_writer.writerow([
                    (time.time()-start_time)/5,
                    current_time_phase,
                    *all_status,  # Unpacks all 4 status values
                    sum(all_status)  # 计算总状态
                ])
                csv_file.flush()  # Ensure data is written to disk
            # 步进仿真
            p.stepSimulation()
            time.sleep(SIMULATION_STEP_TIME)
            
        except KeyboardInterrupt:
            print("仿真线程被中断")
            break
        except Exception as e:
            print(f"仿真出错: {e}")
            break
    
    # 断开PyBullet连接
    if p.isConnected():
        p.disconnect()


def main():
    """主函数"""
    global running
    
    try:
        # 创建四足机器人
        robot = FourLeggedRobot(URDF_FILE)
        
        robot_id, joint_name_to_id, base_link_world_pos_np = robot.init_pybullet(145)# 初始化PyBullet
        leg_chains = robot.create_leg_chains()# 创建运动链
        visualizer = TrajectoryVisualizer()# 创建轨迹可视化器
        # 创建末端执行器标记
        special_markers = visualizer.create_end_effector_markers(leg_chains.keys())
        
        # 重置腿部位置
        robot.reset_leg_positions()
        
        # 创建轨迹规划器
        planner = TrajectoryPlanner()
        
        # 如果启用Tk界面
        if USE_TK_INTERFACE:
            # 创建Tk界面
            ui = TkInterface(robot, planner, visualizer)
            
            # 启动仿真线程
            sim_thread = threading.Thread(
                target=simulation_thread,
                args=(robot, planner, visualizer, ui),
                daemon=True
            )
            sim_thread.start()
            
            # 启动界面
            ui.start()
        else:
            # 不使用Tk界面，直接在主线程中运行仿真
            print("四足机器人仿真已启动，按Ctrl+C停止...")
            simulation_thread(robot, planner, visualizer)
        
    except Exception as e:
        print(f"程序出错: {e}")
    finally:
        # 确保断开PyBullet连接
        if p.isConnected():
            p.disconnect()


if __name__ == "__main__":
    main()
